﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace CTThrottling
{
    public class TokenBucketLimitingService : ILimitingService
    {
        private LimitedQueue<object> _limitedQueue = null;
        private CancellationTokenSource _cancelToken;
        private Task _task = null;
        private int _maxTps;
        private int _limitSize;
        private object _lckObj = new object();
        public TokenBucketLimitingService(int maxTps, int limitSize)
        {
            this._limitSize = limitSize;
            this._maxTps = maxTps;

            if (this._limitSize <= 0)
                this._limitSize = 100;
            if (this._maxTps <= 0)
                this._maxTps = 1;

            _limitedQueue = new LimitedQueue<object>(limitSize);
            for (int i = 0; i < limitSize; i++)
            {
                _limitedQueue.Enqueue(new object());
            }
            _cancelToken = new CancellationTokenSource();
            _task = Task.Factory.StartNew(new Action(TokenProcess), _cancelToken.Token);
        }

        /// <summary>
        /// 定时消息令牌
        /// </summary>
        private void TokenProcess()
        {
            int sleep = 1000 / _maxTps;
            if (sleep == 0)
                sleep = 1;

            DateTime start = DateTime.Now;
            while (_cancelToken.Token.IsCancellationRequested == false)
            {
                try
                {
                    lock (_lckObj)
                    {
                        _limitedQueue.Enqueue(new object());
                    }
                }
                catch
                {
                    //
                }
                finally
                {
                    if (DateTime.Now - start < TimeSpan.FromMilliseconds(sleep))
                    {
                        int newSleep = sleep - (int)(DateTime.Now - start).TotalMilliseconds;
                        if (newSleep > 1)
                            Thread.Sleep(newSleep - 1); //做一下时间上的补偿
                    }
                    start = DateTime.Now;
                }
            }
        }

        /// <summary>
        /// 
        /// </summary>
        public void Dispose()
        {
            _cancelToken.Cancel();
        }

        /// <summary>
        /// 请求令牌
        /// </summary>
        /// <returns>true：获取成功，false：获取失败</returns>
        public bool Request()
        {
            if (_limitedQueue.Count <= 0)
                return false;
            lock (_lckObj)
            {
                if (_limitedQueue.Count <= 0)
                    return false;

                object data = _limitedQueue.Dequeue();
                if (data == null)
                    return false;
            }

            return true;
        }
    }
}
